SEND_RECV = 0
ISEND_IRECV = 1

COLLECTIVE = 2
WAIT_OPS = 3
ALLTOALL = 3
ONETOALL = 4
COMM_OPS = 5

total_proc_num = 16


class mpi_ops() :
    """ 
    describes a single record in the timestamp
    """
    def __init__ (self) :
        self.next_op = []               #the next message on local process
        self.global_coworker = []      # the other calls that collaborates with the current call messege
        self.local_coworker = []       # the wait/waitall messege that on the local process
        self.collective_root = []      # used in bcast and reduce
    def init_value( self, pid, rid,rec) :
        return 

class onetoone (mpi_ops) :
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        self.count = rec[4]
        self.data_type = rec[5]
        self.other = rec[6]    #the other process partisipating the comm
        self.tag = rec[7]
        self.comm = rec[8]
        if ( rec[3] == 2 or rec[3] == 3) :
            self.req = rec[9]
        return 
          
class  wait_ops (mpi_ops) :
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        if ( self.comm_type == 4 ) :
            self.req = rec[4]
        else:
            self.count = rec[4]
            self.reqs = rec[5:] #now it is a list!

class barrier_pos (mpi_ops) :
    def init_value(self, pid, rid, rec):
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        self.comm = rec[4]

class onetoall (mpi_ops) :      #only for bcast (7)
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        self.count = rec[4]
        self.data_type = rec[5]
        self.root = rec[6]
        self.comm = rec[7]

class reduce_ops(mpi_ops) : #for 11 and 12
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        self.count = rec[4]
        self.data_type = rec[5]
        if (self.comm_type == 11 ) :
            self.root = rec[6]
        else :
            self.op = rec[6]
        self.comm = rec[7]

class alltoall (mpi_ops) : #for MPI_Allgather(8) MPI_Alltoall(9), MPI_Alltoallv(10)
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        if ( self.comm_type == 8 or self.comm_type == 9) :
            self.send_cnt = rec[4]
            self.send_type = rec[5]
            self.recv_cnt= rec[6]
            self.revc_type = rec[7]
            self.comm = rec[8]
        elif (self.comm_type == 10 ) :
            self.send_cnt = rec[8]
            self.send_type = rec[4]
            self.recv_cnt= rec[10]
            self.revc_type = rec[5]
            self.comm = rec[6]
            self.list_size = rec[7]
            self.sdispls = rec[11]
            self.rdispls = rec[12]


class comm_ops (mpi_ops) :
    def init_value(self, pid, rid, rec) : #rec is a integer list
        self.pid = pid
        self.rid = rid
        self.statck_id = rec[0]
        self.tp = rec[1]
        self.td = rec[2]
        self.comm_type = rec[3]
        if (self.comm_type == 13) :
            self.comm = rec[4]
            self.comm_out = rec[5]
        elif (self.comm_type == 14) :
            self.comm = rec[4]
            self.color = rec[5]
            self.key = rec[6]
            self.comm_out = rec[7]
            
# xxx_group functions: how to generate links among a set of related processes
# wait ops are ingnored
def send_recv_group( msg_group ) :
    """ the msg_group is a group of message objs """
    group_size = len(msg_group)
    assert (group_size == 2)
    msg_group[0].global_coworker = msg_group[1]
    msg_group[1].global_coworker = msg_group[0]
#end def

def isend_irecv_group (msg_group):
    group_size = len(msg_group)
    assert (group_size == 2)
    msg_group[0].global_coworker .append (msg_group[1])
    msg_group[1].global_coworker .append (msg_group[0])

    req_id = msg_group[0].req
    ops_obj = msg_obj[0]    
    while ( ops_obj) :
        if ( ops_obj.comm_type == 4 and ops_obj.req == req_id ):
            msg_group[0].local_coworker.append(ops_obj)
            break
        elif ( ops_obj.comm_type == 5 and req_id in ops_obj.req) :
            msg_group[0].local_coworker.append(ops_obj)
            break
        ops_obj = ops_obj . next_op
    #end while
    # on the other process
    req_id = msg_group[1].req
    ops_obj = msg_obj[1]    
    while ( ops_obj) :
        if ( ops_obj.comm_type == 4 and ops_obj.req == req_id ):
            msg_group[1].local_coworker.append(ops_obj)
            break
        elif ( ops_obj.comm_type == 5 and req_id in ops_obj.req) :
            msg_group[1].local_coworker.append(ops_obj)
            break
        ops_obj = ops_obj . next_op

    #end while
#end def        


def onetoall_group (msg_group ): # bcast and reduce
    group_size = len(msg_group)
    for i in range(group_size):
        for j in range(group_size) :
            if ( i != j) :
                msg_group[i].global_coworker.append(msg_group[j])
        #end for
    #end for
    root = msg_group[0].root

    for i in range (group_size):
        if ( i != msg_group[i].root ):
            msg_group[root].collective_root.append(msg_group[i]);
    #end for
#end def


def alltoall_group(msg_group) : # barrier, allxxx
    group_size = len(msg_group)
    for i in range(group_size):
        for j in range(group_size) :
            if ( i != j) :
                msg_group[i].global_coworker.append(msg_group[j])
        #end for
    #end for

def comm_ops_group(msg_group): # comm_split comm_dup
    group_size = len(msg_group)
    for i in range(group_size):
        for j in range(group_size) :
            if ( i != j) :
                msg_group[i].global_coworker.append(msg_group[j])
        #end for
    #end for
    


def has_onetoone_grp( curr_msgs, visited_group):
    """ calculates the total number of one to one ops in the current mpi call 
    context, returns the index list and the total number
    for balancing, the ops should not be visited by last time"""
    onetoone_list = []
    onetoone_num = 0   #should have even number
    for i in range(len(curr_msgs )):
        if (curr_msgs[i].comm_type >= 0 and 
            curr_msgs[i].comm_type < 4 and
            visited_group[i] == False) :
            onetoone_num += 1
            onetoone_list .append(i)
    #end for
    assert (onetoone_num%2 == 0)
    return [onetoone_list, onetoone_num]

def extract_grp(curr_msgs, visited_group):
    """ return the group index and the comm_type
    the group indexes are the set of the indexes for each
    call participating the comm, the comm_type are defined at the beginning"""
    ops_group = []
    ret_val = []
    # skip all the wait(all) ops
    for i in range (len (curr_msgs)):
        if ( curr_msgs[i].comm_type == 4 or curr_msgs[i].comm_type == 5) :
            ops_group.append (i)
    if( len (ops_group) ) :
        return [ops_group, WAIT_OPS]

    #check the send/recv and isend/irecv
    [onetoone_list, onetoone_num ] = has_onetoone_grp(curr_msgs, visited_group)
    print "onetoone_num = " + str( onetoone_num)
    if ( onetoone_num > 0) :
        for i in range(len(onetoone_list)) :
            for j in range(i, len(onetoone_list)):
                if (curr_msgs[onetoone_list[j]].other == curr_msgs[onetoone_list[i]].pid and
                    curr_msgs[onetoone_list[j]].pid == curr_msgs[onetoone_list[i]].other and
                    curr_msgs[onetoone_list[j]].tag == curr_msgs[onetoone_list[i]].tag and
                    curr_msgs[onetoone_list[j]].comm == curr_msgs[onetoone_list[i]].comm and
                    (curr_msgs[onetoone_list[j]].comm_type + curr_msgs[onetoone_list[i]].comm_type)%4 == 1):

                    if( curr_msgs[onetoone_list[j]].comm_type <= 1):
                        return [ [onetoone_list[i], onetoone_list[j]], SEND_RECV]
                    else :
                        return [ [onetoone_list[i], onetoone_list[j]], ISEND_IRECV]
        return [None, 0]
    #find the first unvisited call msg

    grp_member_id = total_proc_num

    for i in range(len(curr_msgs)):
        if ( visited_group[i] == 0 ) : #extract the group from here! all collective ops! find all the processes that has the same comm, comm_type
            grp_member_id = i
            break

    if(grp_member_id == total_proc_num ):
        grp_member_id = 0
    #end if

    if(curr_msgs[grp_member_id].comm_type in [6, 8, 9, 10, 12] ): #match comm and comm_type
        for i in range(len(curr_msgs)) :
            if ( curr_msgs[i].comm_type == curr_msgs[grp_member_id].comm_type and curr_msgs[i].comm == curr_msgs[grp_member_id].comm ):
                ops_group.append (i)
        ret_val = [ops_group, ALLTOALL]
    elif( curr_msgs[grp_member_id].comm_type in [7 , 11] ) : #match comm, comm_type and root
        for i in range(len(curr_msgs)): 
            if (  curr_msgs[i].comm_type == curr_msgs[grp_member_id].comm_type and curr_msgs[i].root == curr_msgs[grp_member_id].root and curr_msgs[i].comm == curr_msgs[grp_member_id].comm ):
                ops_group.append (i)
        ret_val = [ops_group, ONETOALL ]
    elif (curr_msgs[grp_member_id].comm_type in [13, 14]) :
        for i in range(len(curr_msgs)):
            if ( curr_msgs[i].comm_type == curr_msgs[grp_member_id].comm_type and curr_msgs[i].comm == curr_msgs[grp_member_id].comm ):
                ops_group.append (i)
        ret_val = [ops_group, COMM_OPS]
    else :
        assert ( False )
    
    return ret_val

            

def link_msg(total_messages) :#generate the partial relationship among these messages
    """ link the messages that work together """
    #in this link algorithm, all the onetoone operations have priority over the other kind of mpi calls

#    curr_msg_ids = [] # the ids of the current on each process

    curr_msgs = []  # this is the current msg objes we need to process
    total_process_num = len(total_messages)
    for i in range(total_process_num) :
        curr_msgs.append(total_messages[i][0])

    #first create the next op link
    for i in range(total_process_num) :
        for msg_id in range(1, len(total_messages[i])):
            total_messages[i][msg_id - 1].next_op = total_messages[i][msg_id]
        #end for
    #endfor

    counter = 0
    visited_group = [0, 0, 0,0, 0, 0,0, 0, 0,0, 0, 0,0, 0, 0,0]  #used to make sure the extracted group are balanced
    while (True ) :
        for i in range(total_proc_num):
            print "(" + str(curr_msgs[i].comm_type) + " " + str(visited_group[i]) + " " + str (curr_msgs[i].comm) + ")" ,
        print ""

        

        [grp_index, comm_type ] = extract_grp(curr_msgs, visited_group)
        assert (grp_index )
        if (comm_type == WAIT_OPS ) : #do nothing
            pass
        elif comm_type == SEND_RECV :
            send_recv_group([ curr_msgs[j]  for j in grp_index ] )
        elif comm_type == ISEND_IRECV :
            isend_irecv_group([ curr_msgs[j] for j in grp_index ] )
        elif comm_type == ALLTOALL :
            alltoall_group([ curr_msgs[j]  for j in grp_index ] )
        elif comm_type == ONETOALL :
            onetoall_group([ curr_msgs[j]  for j in grp_index ] )
        elif comm_type == COMM_OPS :
            comm_ops_group([ curr_msgs[j]  for j in grp_index ] )

        visited_group = [ 0 for i in range (total_proc_num)] # show which column is visited in the previous visit
        for i in grp_index :
            curr_msgs[i] = curr_msgs[i].next_op

        if len(grp_index) < total_proc_num :
            for i in grp_index :
                visited_group[i] = 1

        counter += 1
        if (counter > total_proc_num * 400) :
            break
# end



def show_objs(total_messages):
    """
    print the objects
    """
    for process_msg in total_messages :
        for msg in process_msg :
            print msg.__dict__
#end of the show_objs function


def process_rec( line, pid, rid) :
    """
    process the line!
    """

    rec_list = line.split(",")
    new_list = []
    for var in rec_list :
        new_list.append( int( var.strip()))
    
    rec_list = new_list

    if (rec_list[3] < 4 ):
        new_msg = onetoone()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] < 6) :
        new_msg = wait_ops()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] == 6 ) :
        new_msg = barrier_pos()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] == 7):
        new_msg = onetoall()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] <= 10):
        new_msg = alltoall()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] <= 12):
        new_msg = reduce_ops()
        new_msg.init_value(pid, rid, rec_list)
    elif (rec_list[3] <= 14):
        new_msg = comm_ops()
        new_msg.init_value(pid, rid, rec_list)
        
    return new_msg
#end of def process_rec


def read_data():
    """
    parse the input files for traces, and put the data into the objects
    """
    print "start processing!"
    total_messages = [] #containing all the messages from all processes, arranged as the list
    for i in range(total_proc_num) : 
        if ( i < 10 ) :
            file_name = "000000" + str(i)
        else:
            file_name = "00000" + str(i)
        file = open("prof/bt.C.16/" + file_name)
        if (not file ):
            print "open file failed!"

        pid = i
        seen_rec = 0
       	process_messages = [] # the messages per process


        for line in file :
            if (line.startswith("#") or line.startswith("reqs:") or line.startswith("**")):
                pass
            else :
                message = process_rec(line, pid, seen_rec)
                seen_rec += 1
#                print type(message), message.__class__.__name__
                
                process_messages.append(message)
                if (seen_rec >=400 ):
                    break
        #end of for loop
        total_messages.append(process_messages)
        file.close()
    #end loop
    return total_messages
#end of read_data


if __name__ == "__main__" :
    total_messages = read_data()
 #   i = 0
#    for process_msg in total_messages :
 #       print "start of the message at process " + str(i)
  #   for msg in process_msg:
     #       if( hasattr(msg, "comm")) :
      #          print msg.comm
       # i = i + 1
    #end for
    link_msg(total_messages)
